import os
import joblib
import numpy as np
from PIL import Image
import gradio as gr

# 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
try:
    base_dir = os.path.dirname(__file__)
except NameError:
    # 在交互式环境（Jupyter/REPL）中 __file__ 可能未定义，使用当前工作目录作为回退
    base_dir = os.getcwd()
model_path = os.path.join(base_dir, "best_knn_model.pkl")
if not os.path.exists(model_path):
    raise FileNotFoundError(f"找不到模型文件: {model_path}，请先运行 optimal_knn.py 以训练并保存模型。")
knn = joblib.load(model_path)

# 定义预测函数，这个函数将用于Gradio接口进行预测
def preprocess_image(img):
    if img is None:
        return None
    # img 可能是 (H,W) 或 (H,W,3)
    if img.ndim == 3:
        img = np.mean(img, axis=2)
    pil = Image.fromarray(img.astype(np.uint8)).convert("L")
    pil = pil.resize((8, 8), Image.ANTIALIAS)
    arr = np.asarray(pil, dtype=np.float32)
    # 将 0-255 缩放到 0-16（与 sklearn.datasets.load_digits 的范围一致）
    arr = (arr / 255.0) * 16.0
    return arr.reshape(1, -1)

def predict(img):
    x = preprocess_image(img)
    if x is None:
        return {}, {}
    if hasattr(knn, "predict_proba"):
        probs = knn.predict_proba(x)[0]
    else:
        pred = int(knn.predict(x)[0])
        probs = np.zeros(10, dtype=float)
        probs[pred] = 1.0
    # 构造标签->概率的字典（字符串键）
    prob_dict = {str(i): float(probs[i]) for i in range(len(probs))}
    # 返回两个输出：第一个用于显示 top3，第二个用于显示完整置信度字典
    return prob_dict, prob_dict

# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", image_mode="L"),
    outputs=[
        gr.Label(num_top_classes=3, label="预测（Top-3）"),
        gr.Label(label="置信度（每个类别）")
    ],
    title="KNN digits 预测",
    description="上传手写数字图像，界面会把图像缩放到8x8并使用保存的KNN模型进行预测。"
)

# 启动Gradio接口，用户可以通过这个接口进行交互
if __name__ == "__main__":
    iface.launch()
